"use client";

import * as React from "react";
import TextareaAutosize, {
  type TextareaAutosizeProps,
} from "react-textarea-autosize";

import { type TEquationElement } from "platejs";
import { type PlateElementProps } from "platejs/react";

import { useEquationElement, useEquationInput } from "@platejs/math/react";
import { BlockSelectionPlugin } from "@platejs/selection/react";
import { CornerDownLeftIcon, RadicalIcon } from "lucide-react";
import {
  createPrimitiveComponent,
  PlateElement,
  useEditorRef,
  useEditorSelector,
  useElement,
  useReadOnly,
  useSelected,
} from "platejs/react";

import { Button } from "@/components/plate/ui/button";
import {
  Popover,
  PopoverContent,
  PopoverTrigger,
} from "@/components/plate/ui/popover";
import { cn } from "@/lib/utils";

export function EquationElement(props: PlateElementProps&lt;TEquationElement&gt;) {
  const selected = useSelected();
  const [open, setOpen] = React.useState(selected);
  const katexRef = React.useRef&lt;HTMLDivElement | null&gt;(null);

  useEquationElement({
    element: props.element,
    katexRef: katexRef,
    options: {
      displayMode: true,
      errorColor: "#cc0000",
      fleqn: false,
      leqno: false,
      macros: { "\\f": "#1f(#2)" },
      output: "htmlAndMathml",
      strict: "warn",
      throwOnError: false,
      trust: false,
    },
  });

  return (
    &lt;PlateElement className="my-1" {...props}&gt;
      &lt;Popover open={open} onOpenChange={setOpen} modal={false}&gt;
        &lt;PopoverTrigger asChild&gt;
          <div data-selected="true]:bg-primary/10&quot;," data-selected="{selected}">
            {props.element.texExpression.length &gt; 0 ? (
              <span />
            ) : (
              <div>
                &lt;RadicalIcon className="size-6 text-muted-foreground/80" /&gt;
                <div>Add a Tex equation</div>
              </div>
            )}
          </div>
        &lt;/PopoverTrigger&gt;

        &lt;EquationPopoverContent
          open={open}
          placeholder={`f(x) = \\begin{cases}\n  x^2, &\\quad x &gt; 0 \\\\\n  0, &\\quad x = 0 \\\\\n  -x^2, &\\quad x &lt; 0\n\\end{cases}`}
          isInline={false}
          setOpen={setOpen}
        /&gt;
      &lt;/Popover&gt;

      {props.children}
    &lt;/PlateElement&gt;
  );
}

export function InlineEquationElement(
  props: PlateElementProps&lt;TEquationElement&gt;,
) {
  const element = props.element;
  const katexRef = React.useRef&lt;HTMLDivElement | null&gt;(null);
  const selected = useSelected();
  const isCollapsed = useEditorSelector(
    (editor) =&gt; editor.api.isCollapsed(),
    [],
  );
  const [open, setOpen] = React.useState(selected && isCollapsed);

  React.useEffect(() =&gt; {
    if (selected && isCollapsed) {
      setOpen(true);
    }
  }, [selected, isCollapsed]);

  useEquationElement({
    element,
    katexRef: katexRef,
    options: {
      displayMode: true,
      errorColor: "#cc0000",
      fleqn: false,
      leqno: false,
      macros: { "\\f": "#1f(#2)" },
      output: "htmlAndMathml",
      strict: "warn",
      throwOnError: false,
      trust: false,
    },
  });

  return (
    &lt;PlateElement
      {...props}
      className={cn(
        "[&_.katex-display]:my-0! mx-1 inline-block select-none rounded-sm",
      )}
    &gt;
      &lt;Popover open={open} onOpenChange={setOpen} modal={false}&gt;
        &lt;PopoverTrigger asChild&gt;
          <div> 0 && open) || selected) &&
                "after:bg-brand/15",
              element.texExpression.length === 0 &&
                "text-muted-foreground after:bg-neutral-500/10",
            )}
            contentEditable={false}
          &gt;
            <span />
            {element.texExpression.length === 0 && (
              <span>
                &lt;RadicalIcon className="mr-1 inline-block h-[19px] w-4 py-[1.5px] align-text-bottom" /&gt;
                New equation
              </span>
            )}
          </div>
        &lt;/PopoverTrigger&gt;

        &lt;EquationPopoverContent
          className="my-auto"
          open={open}
          placeholder="E = mc^2"
          setOpen={setOpen}
          isInline
        /&gt;
      &lt;/Popover&gt;

      {props.children}
    &lt;/PlateElement&gt;
  );
}

const EquationInput = createPrimitiveComponent(TextareaAutosize)({
  propsHook: useEquationInput,
});

const EquationPopoverContent = ({
  className,
  isInline,
  open,
  setOpen,
  ...props
}: {
  isInline: boolean;
  open: boolean;
  setOpen: (open: boolean) =&gt; void;
} & TextareaAutosizeProps) =&gt; {
  const editor = useEditorRef();
  const readOnly = useReadOnly();
  const element = useElement&lt;TEquationElement&gt;();

  React.useEffect(() =&gt; {
    if (isInline && open) {
      setOpen(true);
    }
  }, [isInline, open, setOpen]);

  if (readOnly) return null;

  const onClose = () =&gt; {
    setOpen(false);

    if (isInline) {
      editor.tf.select(element, { focus: true, next: true });
    } else {
      editor
        .getApi(BlockSelectionPlugin)
        .blockSelection.set(element.id as string);
    }
  };

  return (
    &lt;PopoverContent
      className="flex gap-2"
      onEscapeKeyDown={(e) =&gt; {
        e.preventDefault();
      }}
      contentEditable={false}
    &gt;
      &lt;EquationInput
        className={cn("max-h-[50vh] grow resize-none p-2 text-sm", className)}
        state={{ isInline, open, onClose }}
        autoFocus
        {...props}
      /&gt;

      &lt;Button variant="secondary" className="px-3" onClick={onClose}&gt;
        Done &lt;CornerDownLeftIcon className="size-3.5" /&gt;
      &lt;/Button&gt;
    &lt;/PopoverContent&gt;
  );
};
